{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import copy\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import torch.backends.cudnn as cudnn\n",
    "from nets.yolo4 import YoloBody\n",
    "from nets.yolo_training import YOLOLoss, Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#---------------------------------------------------#\n",
    "#   获得类和先验框\n",
    "#---------------------------------------------------#\n",
    "def get_classes(classes_path):\n",
    "    '''loads the classes'''\n",
    "    with open(classes_path) as f:\n",
    "        class_names = f.readlines()\n",
    "    class_names = [c.strip() for c in class_names]\n",
    "    return class_names\n",
    "\n",
    "\n",
    "def get_anchors(anchors_path):\n",
    "    '''loads the anchors from a file'''\n",
    "    with open(anchors_path) as f:\n",
    "        anchors = f.readline()\n",
    "    anchors = [float(x) for x in anchors.split(',')]\n",
    "    return np.array(anchors).reshape([-1,3,2])[::-1,:,:]\n",
    "\n",
    "\n",
    "#---------------------------------------------------#\n",
    "#   训练一个epoch\n",
    "#---------------------------------------------------#\n",
    "def fit_one_epoch(net, yolo_losses, epoch, epoch_size, epoch_size_val, gen,genval, Epoch, cuda, optimizer, lr_scheduler):\n",
    "    total_loss = 0\n",
    "    val_loss = 0\n",
    "    print('\\n' + '-' * 10 + 'Train one epoch.' + '-' * 10)\n",
    "    print('Epoch:'+ str(epoch+1) + '/' + str(Epoch))\n",
    "    print('Start Training.')\n",
    "    net.train()\n",
    "    for iteration in range(epoch_size):\n",
    "        start_time = time.time()\n",
    "        images, targets = next(gen)\n",
    "        with torch.no_grad():\n",
    "            if cuda:\n",
    "                images = Variable(torch.from_numpy(images).type(torch.FloatTensor)).cuda()\n",
    "                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]\n",
    "            else:\n",
    "                images = Variable(torch.from_numpy(images).type(torch.FloatTensor))\n",
    "                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(images)\n",
    "        losses = []\n",
    "        for i in range(3):\n",
    "            loss_item = yolo_losses[i](outputs[i], targets)\n",
    "            losses.append(loss_item[0])\n",
    "        loss = sum(losses)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        total_loss += loss\n",
    "        waste_time = time.time() - start_time\n",
    "        if iteration == 0 or (iteration+1) % 10 == 0:\n",
    "            print('step:' + str(iteration+1) + '/' + str(epoch_size) + ' || Total Loss: %.4f || %.4fs/step' % (total_loss/(iteration+1), waste_time))\n",
    "    print('Finish Training.')\n",
    "    '''        \n",
    "    print('Start Validation.')\n",
    "    net.eval()\n",
    "    for iteration in range(epoch_size_val):\n",
    "        images_val, targets_val = next(genval)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            if cuda:\n",
    "                images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor)).cuda()\n",
    "                targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]\n",
    "            else:\n",
    "                images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor))\n",
    "                targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]\n",
    "            optimizer.zero_grad()\n",
    "            outputs = net(images_val)\n",
    "            losses = []\n",
    "            for i in range(3):\n",
    "                loss_item = yolo_losses[i](outputs[i], targets_val)\n",
    "                losses.append(loss_item[0])\n",
    "            loss = sum(losses)\n",
    "            val_loss += loss\n",
    "    print('Finish Validation')\n",
    "    '''\n",
    "    print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss/(epoch_size+1), val_loss/(epoch_size_val+1)))\n",
    "    \n",
    "    return total_loss/(epoch_size+1), val_loss/(epoch_size_val+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-------------------------------#\n",
    "#   输入的shape大小\n",
    "#   显存比较小可以使用416x416\n",
    "#   显存比较大可以使用608x608\n",
    "#-------------------------------#\n",
    "#input_shape = (416,416)\n",
    "input_shape = (608, 608)\n",
    "\n",
    "#-------------------------------#\n",
    "#   tricks的使用设置\n",
    "#-------------------------------#\n",
    "Cosine_lr = True\n",
    "mosaic = True\n",
    "# 用于设定是否使用cuda\n",
    "Cuda = True\n",
    "smoooth_label = 0.03\n",
    "\n",
    "#-------------------------------#\n",
    "#   获得训练集和验证集的annotations\n",
    "#   \n",
    "#-------------------------------#\n",
    "train_annotation_path = 'model_data/mask_train.txt'\n",
    "val_annotation_path = 'model_data/mask_val.txt'\n",
    "\n",
    "#-------------------------------#\n",
    "#   获得先验框和类\n",
    "#-------------------------------#\n",
    "anchors_path = 'model_data/yolo_anchors.txt'\n",
    "classes_path = 'model_data/mask_classes.txt'   \n",
    "class_names = get_classes(classes_path)\n",
    "anchors = get_anchors(anchors_path)\n",
    "num_classes = len(class_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained model weights.\n",
      "Finished!\n"
     ]
    }
   ],
   "source": [
    "# 创建模型\n",
    "model = YoloBody(len(anchors[0]), num_classes)\n",
    "#model_path = \"model_data/yolov4_coco_pretrained_weights.pth\"\n",
    "model_path = \"model_data/yolov4_maskdetect_weights0.pth\"\n",
    "# 加快模型训练的效率\n",
    "print('Loading pretrained model weights.')\n",
    "model_dict = model.state_dict()\n",
    "pretrained_dict = torch.load(model_path)\n",
    "pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}\n",
    "model_dict.update(pretrained_dict)\n",
    "model.load_state_dict(model_dict)\n",
    "print('Finished!')\n",
    "\n",
    "if Cuda:\n",
    "    net = torch.nn.DataParallel(model)\n",
    "    cudnn.benchmark = True\n",
    "    net = net.cuda()\n",
    "\n",
    "# 建立loss函数\n",
    "yolo_losses = []\n",
    "for i in range(3):\n",
    "    yolo_losses.append(YOLOLoss(np.reshape(anchors, [-1,2]), num_classes, \\\n",
    "                                (input_shape[1], input_shape[0]), smoooth_label, Cuda))\n",
    "# read train lines and val lines\n",
    "with open(train_annotation_path) as f:\n",
    "    train_lines = f.readlines()\n",
    "with open(val_annotation_path) as f:\n",
    "    val_lines = f.readlines()\n",
    "num_train = len(train_lines)\n",
    "num_val = len(val_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:1/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 15688.8691 || 6.8430s/step\n",
      "step:10/120 || Total Loss: 12204.3193 || 1.6851s/step\n",
      "step:20/120 || Total Loss: 9636.8506 || 1.4570s/step\n",
      "step:30/120 || Total Loss: 7773.7681 || 1.6354s/step\n",
      "step:40/120 || Total Loss: 6444.1396 || 1.7304s/step\n",
      "step:50/120 || Total Loss: 5482.1118 || 1.5075s/step\n",
      "step:60/120 || Total Loss: 4764.4355 || 1.3454s/step\n",
      "step:70/120 || Total Loss: 4211.4189 || 1.4012s/step\n",
      "step:80/120 || Total Loss: 3774.4348 || 1.7187s/step\n",
      "step:90/120 || Total Loss: 3421.0762 || 1.5336s/step\n",
      "step:100/120 || Total Loss: 3128.3455 || 1.4144s/step\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/myconda/lib/python3.7/site-packages/PIL/TiffImagePlugin.py:804: UserWarning: Corrupt EXIF data.  Expecting to read 4 bytes but only got 0. \n",
      "  warnings.warn(str(msg))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step:110/120 || Total Loss: 2882.8147 || 1.3798s/step\n",
      "step:120/120 || Total Loss: 2674.7305 || 1.4065s/step\n",
      "Finish Training.\n",
      "Total Loss: 2652.6250 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:2/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 359.2517 || 1.8192s/step\n",
      "step:10/120 || Total Loss: 345.9673 || 1.3860s/step\n",
      "step:20/120 || Total Loss: 328.0730 || 2.3537s/step\n",
      "step:30/120 || Total Loss: 319.5523 || 1.9720s/step\n",
      "step:40/120 || Total Loss: 301.3265 || 1.2407s/step\n",
      "step:50/120 || Total Loss: 288.6107 || 1.5783s/step\n",
      "step:60/120 || Total Loss: 278.3066 || 1.6308s/step\n",
      "step:70/120 || Total Loss: 267.8594 || 1.6753s/step\n",
      "step:80/120 || Total Loss: 257.0878 || 1.4918s/step\n",
      "step:90/120 || Total Loss: 248.1550 || 1.1582s/step\n",
      "step:100/120 || Total Loss: 239.4298 || 1.2975s/step\n",
      "step:110/120 || Total Loss: 235.7772 || 1.4822s/step\n",
      "step:120/120 || Total Loss: 228.1834 || 1.4821s/step\n",
      "Finish Training.\n",
      "Total Loss: 226.2975 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:3/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 124.2823 || 1.2445s/step\n",
      "step:10/120 || Total Loss: 164.8357 || 2.1981s/step\n",
      "step:20/120 || Total Loss: 149.4365 || 1.4476s/step\n",
      "step:30/120 || Total Loss: 152.2224 || 1.8367s/step\n",
      "step:40/120 || Total Loss: 157.9434 || 1.3714s/step\n",
      "step:50/120 || Total Loss: 153.4084 || 1.3691s/step\n",
      "step:60/120 || Total Loss: 150.1354 || 1.5633s/step\n",
      "step:70/120 || Total Loss: 143.9465 || 1.7103s/step\n",
      "step:80/120 || Total Loss: 138.9583 || 1.4637s/step\n",
      "step:90/120 || Total Loss: 134.4086 || 1.2036s/step\n",
      "step:100/120 || Total Loss: 131.8153 || 1.5356s/step\n",
      "step:110/120 || Total Loss: 127.6153 || 1.4286s/step\n",
      "step:120/120 || Total Loss: 125.0696 || 1.2749s/step\n",
      "Finish Training.\n",
      "Total Loss: 124.0360 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:4/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 127.6602 || 1.7243s/step\n",
      "step:10/120 || Total Loss: 101.3381 || 1.6267s/step\n",
      "step:20/120 || Total Loss: 97.2072 || 1.5298s/step\n",
      "step:30/120 || Total Loss: 101.5166 || 1.6919s/step\n",
      "step:40/120 || Total Loss: 95.8324 || 1.4091s/step\n",
      "step:50/120 || Total Loss: 91.6894 || 1.6408s/step\n",
      "step:60/120 || Total Loss: 91.0613 || 1.3112s/step\n",
      "step:70/120 || Total Loss: 90.5509 || 1.8081s/step\n",
      "step:80/120 || Total Loss: 92.6823 || 1.7642s/step\n",
      "step:90/120 || Total Loss: 94.1193 || 1.4574s/step\n",
      "step:100/120 || Total Loss: 92.2101 || 1.7817s/step\n",
      "step:110/120 || Total Loss: 90.2761 || 1.8930s/step\n",
      "step:120/120 || Total Loss: 88.5604 || 1.3168s/step\n",
      "Finish Training.\n",
      "Total Loss: 87.8285 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:5/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 69.0592 || 1.5089s/step\n",
      "step:10/120 || Total Loss: 96.6841 || 1.4639s/step\n",
      "step:20/120 || Total Loss: 78.2830 || 2.1573s/step\n",
      "step:30/120 || Total Loss: 69.9976 || 1.3906s/step\n",
      "step:40/120 || Total Loss: 66.3320 || 1.7036s/step\n",
      "step:50/120 || Total Loss: 69.0854 || 1.3701s/step\n",
      "step:60/120 || Total Loss: 68.1419 || 2.0508s/step\n",
      "step:70/120 || Total Loss: 69.5947 || 1.7939s/step\n",
      "step:80/120 || Total Loss: 68.0882 || 1.3704s/step\n",
      "step:90/120 || Total Loss: 65.6736 || 1.2152s/step\n",
      "step:100/120 || Total Loss: 64.9521 || 1.4611s/step\n",
      "step:110/120 || Total Loss: 65.5481 || 1.6082s/step\n",
      "step:120/120 || Total Loss: 64.2527 || 1.5804s/step\n",
      "Finish Training.\n",
      "Total Loss: 63.7217 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:6/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 50.4700 || 1.2556s/step\n",
      "step:10/120 || Total Loss: 52.1577 || 1.4906s/step\n",
      "step:20/120 || Total Loss: 55.8429 || 1.6901s/step\n",
      "step:30/120 || Total Loss: 72.8140 || 1.7906s/step\n",
      "step:40/120 || Total Loss: 66.9406 || 1.7587s/step\n",
      "step:50/120 || Total Loss: 70.1888 || 1.5041s/step\n",
      "step:60/120 || Total Loss: 67.9102 || 1.4900s/step\n",
      "step:70/120 || Total Loss: 65.0673 || 1.5760s/step\n",
      "step:80/120 || Total Loss: 61.7709 || 1.4417s/step\n",
      "step:90/120 || Total Loss: 60.8019 || 1.6868s/step\n",
      "step:100/120 || Total Loss: 58.6333 || 1.3259s/step\n",
      "step:110/120 || Total Loss: 58.1179 || 1.9466s/step\n",
      "step:120/120 || Total Loss: 57.8835 || 1.7298s/step\n",
      "Finish Training.\n",
      "Total Loss: 57.4051 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:7/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 24.9863 || 1.4087s/step\n",
      "step:10/120 || Total Loss: 64.8083 || 1.6963s/step\n",
      "step:20/120 || Total Loss: 54.8125 || 1.2989s/step\n",
      "step:30/120 || Total Loss: 48.9344 || 1.7102s/step\n",
      "step:40/120 || Total Loss: 47.1486 || 1.5859s/step\n",
      "step:50/120 || Total Loss: 46.8563 || 1.8071s/step\n",
      "step:60/120 || Total Loss: 49.7594 || 1.7835s/step\n",
      "step:70/120 || Total Loss: 50.6075 || 2.0429s/step\n",
      "step:80/120 || Total Loss: 50.3934 || 1.5747s/step\n",
      "step:90/120 || Total Loss: 48.8376 || 1.3688s/step\n",
      "step:100/120 || Total Loss: 47.9427 || 1.3055s/step\n",
      "step:110/120 || Total Loss: 47.2428 || 1.5705s/step\n",
      "step:120/120 || Total Loss: 46.7030 || 1.4881s/step\n",
      "Finish Training.\n",
      "Total Loss: 46.3170 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:8/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 83.1002 || 1.6849s/step\n",
      "step:10/120 || Total Loss: 43.2236 || 1.4533s/step\n",
      "step:20/120 || Total Loss: 37.0472 || 1.4180s/step\n",
      "step:30/120 || Total Loss: 50.0718 || 1.8619s/step\n",
      "step:40/120 || Total Loss: 48.0531 || 1.3468s/step\n",
      "step:50/120 || Total Loss: 45.0169 || 1.6498s/step\n",
      "step:60/120 || Total Loss: 42.4451 || 1.7557s/step\n",
      "step:70/120 || Total Loss: 40.9353 || 1.6009s/step\n",
      "step:80/120 || Total Loss: 40.0794 || 1.5302s/step\n",
      "step:90/120 || Total Loss: 39.7884 || 1.4759s/step\n",
      "step:100/120 || Total Loss: 40.9428 || 1.3096s/step\n",
      "step:110/120 || Total Loss: 40.7345 || 2.4208s/step\n",
      "step:120/120 || Total Loss: 41.6421 || 1.5526s/step\n",
      "Finish Training.\n",
      "Total Loss: 41.2980 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:9/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 30.6533 || 1.7020s/step\n",
      "step:10/120 || Total Loss: 36.7151 || 1.4711s/step\n",
      "step:20/120 || Total Loss: 34.9089 || 1.6169s/step\n",
      "step:30/120 || Total Loss: 37.8552 || 2.0150s/step\n",
      "step:40/120 || Total Loss: 36.3650 || 2.0132s/step\n",
      "step:50/120 || Total Loss: 36.4452 || 1.9392s/step\n",
      "step:60/120 || Total Loss: 39.3719 || 1.8061s/step\n",
      "step:70/120 || Total Loss: 37.5438 || 1.7135s/step\n",
      "step:80/120 || Total Loss: 38.3320 || 1.8169s/step\n",
      "step:90/120 || Total Loss: 37.9111 || 1.3103s/step\n",
      "step:100/120 || Total Loss: 37.6068 || 1.6984s/step\n",
      "step:110/120 || Total Loss: 37.2881 || 1.2603s/step\n",
      "step:120/120 || Total Loss: 37.0020 || 1.5437s/step\n",
      "Finish Training.\n",
      "Total Loss: 36.6962 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:10/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 28.8127 || 1.5477s/step\n",
      "step:10/120 || Total Loss: 35.7997 || 1.4599s/step\n",
      "step:20/120 || Total Loss: 36.0810 || 1.2531s/step\n",
      "step:30/120 || Total Loss: 35.6410 || 1.5023s/step\n",
      "step:40/120 || Total Loss: 34.9987 || 1.5378s/step\n",
      "step:50/120 || Total Loss: 33.5137 || 1.5870s/step\n",
      "step:60/120 || Total Loss: 35.0475 || 1.4981s/step\n",
      "step:70/120 || Total Loss: 34.3198 || 1.8421s/step\n",
      "step:80/120 || Total Loss: 35.7354 || 1.8494s/step\n",
      "step:90/120 || Total Loss: 34.4763 || 1.2467s/step\n",
      "step:100/120 || Total Loss: 34.2314 || 1.7309s/step\n",
      "step:110/120 || Total Loss: 34.3495 || 1.3269s/step\n",
      "step:120/120 || Total Loss: 34.0682 || 1.4525s/step\n",
      "Finish Training.\n",
      "Total Loss: 33.7867 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:11/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 21.0243 || 1.6315s/step\n",
      "step:10/120 || Total Loss: 47.0997 || 1.2918s/step\n",
      "step:20/120 || Total Loss: 34.8475 || 2.1431s/step\n",
      "step:30/120 || Total Loss: 32.5683 || 1.6228s/step\n",
      "step:40/120 || Total Loss: 30.0201 || 1.1105s/step\n",
      "step:50/120 || Total Loss: 30.3979 || 1.6447s/step\n",
      "step:60/120 || Total Loss: 30.6778 || 1.9496s/step\n",
      "step:70/120 || Total Loss: 31.0141 || 1.4351s/step\n",
      "step:80/120 || Total Loss: 31.2245 || 1.4105s/step\n",
      "step:90/120 || Total Loss: 31.6073 || 1.8010s/step\n",
      "step:100/120 || Total Loss: 31.8341 || 1.5935s/step\n",
      "step:110/120 || Total Loss: 31.6581 || 1.4046s/step\n",
      "step:120/120 || Total Loss: 32.7015 || 1.5503s/step\n",
      "Finish Training.\n",
      "Total Loss: 32.4313 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:12/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 14.4423 || 1.4980s/step\n",
      "step:10/120 || Total Loss: 22.3417 || 1.6391s/step\n",
      "step:20/120 || Total Loss: 42.4041 || 1.7587s/step\n",
      "step:30/120 || Total Loss: 36.0802 || 1.6885s/step\n",
      "step:40/120 || Total Loss: 34.8042 || 1.7496s/step\n",
      "step:50/120 || Total Loss: 32.9821 || 1.6352s/step\n",
      "step:60/120 || Total Loss: 32.2740 || 1.3920s/step\n",
      "step:70/120 || Total Loss: 30.9345 || 1.4033s/step\n",
      "step:80/120 || Total Loss: 30.8213 || 2.2981s/step\n",
      "step:90/120 || Total Loss: 29.6906 || 1.6875s/step\n",
      "step:100/120 || Total Loss: 30.5643 || 1.6684s/step\n",
      "step:110/120 || Total Loss: 29.8639 || 1.8821s/step\n",
      "step:120/120 || Total Loss: 29.9712 || 1.4804s/step\n",
      "Finish Training.\n",
      "Total Loss: 29.7235 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:13/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 12.5914 || 1.9088s/step\n",
      "step:10/120 || Total Loss: 23.5895 || 1.3081s/step\n",
      "step:20/120 || Total Loss: 23.8874 || 1.1832s/step\n",
      "step:30/120 || Total Loss: 21.9874 || 1.5422s/step\n",
      "step:40/120 || Total Loss: 26.1549 || 1.6603s/step\n",
      "step:50/120 || Total Loss: 25.3819 || 0.9712s/step\n",
      "step:60/120 || Total Loss: 25.3024 || 1.3568s/step\n",
      "step:70/120 || Total Loss: 25.4105 || 1.4280s/step\n",
      "step:80/120 || Total Loss: 25.6688 || 1.7038s/step\n",
      "step:90/120 || Total Loss: 24.8554 || 1.5563s/step\n",
      "step:100/120 || Total Loss: 24.9389 || 1.3715s/step\n",
      "step:110/120 || Total Loss: 25.3168 || 1.5615s/step\n",
      "step:120/120 || Total Loss: 26.0490 || 2.1625s/step\n",
      "Finish Training.\n",
      "Total Loss: 25.8338 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:14/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 48.1161 || 1.5927s/step\n",
      "step:10/120 || Total Loss: 30.5352 || 1.4875s/step\n",
      "step:20/120 || Total Loss: 28.4623 || 1.6631s/step\n",
      "step:30/120 || Total Loss: 29.1715 || 1.6138s/step\n",
      "step:40/120 || Total Loss: 30.2424 || 1.8887s/step\n",
      "step:50/120 || Total Loss: 31.3752 || 1.3834s/step\n",
      "step:60/120 || Total Loss: 29.1351 || 1.7105s/step\n",
      "step:70/120 || Total Loss: 29.8731 || 1.7114s/step\n",
      "step:80/120 || Total Loss: 31.1062 || 1.5580s/step\n",
      "step:90/120 || Total Loss: 30.4017 || 1.6217s/step\n",
      "step:100/120 || Total Loss: 30.7051 || 1.8213s/step\n",
      "step:110/120 || Total Loss: 30.1312 || 1.2234s/step\n",
      "step:120/120 || Total Loss: 29.4390 || 1.4589s/step\n",
      "Finish Training.\n",
      "Total Loss: 29.1957 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:15/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 15.0875 || 1.5941s/step\n",
      "step:10/120 || Total Loss: 25.3374 || 1.1807s/step\n",
      "step:20/120 || Total Loss: 27.7153 || 1.5931s/step\n",
      "step:30/120 || Total Loss: 32.5038 || 1.7127s/step\n",
      "step:40/120 || Total Loss: 30.6845 || 1.6447s/step\n",
      "step:50/120 || Total Loss: 29.9032 || 1.6452s/step\n",
      "step:60/120 || Total Loss: 29.2921 || 1.2295s/step\n",
      "step:70/120 || Total Loss: 28.6460 || 1.3227s/step\n",
      "step:80/120 || Total Loss: 26.8682 || 1.5877s/step\n",
      "step:90/120 || Total Loss: 26.7706 || 1.7155s/step\n",
      "step:100/120 || Total Loss: 26.0471 || 1.4454s/step\n",
      "step:110/120 || Total Loss: 25.7004 || 1.7442s/step\n",
      "step:120/120 || Total Loss: 25.4120 || 1.4313s/step\n",
      "Finish Training.\n",
      "Total Loss: 25.2019 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:16/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 11.8339 || 1.4096s/step\n",
      "step:10/120 || Total Loss: 30.3076 || 1.5929s/step\n",
      "step:20/120 || Total Loss: 26.2466 || 1.2673s/step\n",
      "step:30/120 || Total Loss: 24.5809 || 1.5180s/step\n",
      "step:40/120 || Total Loss: 23.1425 || 1.5836s/step\n",
      "step:50/120 || Total Loss: 25.6440 || 1.5080s/step\n",
      "step:60/120 || Total Loss: 24.9481 || 1.9047s/step\n",
      "step:70/120 || Total Loss: 25.7763 || 1.6877s/step\n",
      "step:80/120 || Total Loss: 26.3741 || 1.3121s/step\n",
      "step:90/120 || Total Loss: 25.9405 || 1.3853s/step\n",
      "step:100/120 || Total Loss: 25.6240 || 1.4464s/step\n",
      "step:110/120 || Total Loss: 25.7630 || 2.0382s/step\n",
      "step:120/120 || Total Loss: 25.1634 || 1.5205s/step\n",
      "Finish Training.\n",
      "Total Loss: 24.9555 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:17/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 16.8495 || 1.4896s/step\n",
      "step:10/120 || Total Loss: 22.7047 || 1.9003s/step\n",
      "step:20/120 || Total Loss: 53.0835 || 2.7320s/step\n",
      "step:30/120 || Total Loss: 58.3232 || 1.6838s/step\n",
      "step:40/120 || Total Loss: 52.4184 || 1.4265s/step\n",
      "step:50/120 || Total Loss: 47.1840 || 1.6752s/step\n",
      "step:60/120 || Total Loss: 43.4517 || 1.4905s/step\n",
      "step:70/120 || Total Loss: 40.4032 || 1.5979s/step\n",
      "step:80/120 || Total Loss: 38.1252 || 1.6288s/step\n",
      "step:90/120 || Total Loss: 35.9339 || 2.2201s/step\n",
      "step:100/120 || Total Loss: 34.1183 || 1.6011s/step\n",
      "step:110/120 || Total Loss: 32.7905 || 1.5034s/step\n",
      "step:120/120 || Total Loss: 32.2384 || 1.3709s/step\n",
      "Finish Training.\n",
      "Total Loss: 31.9719 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:18/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 6.2830 || 1.5172s/step\n",
      "step:10/120 || Total Loss: 32.8177 || 1.1460s/step\n",
      "step:20/120 || Total Loss: 26.9262 || 1.6132s/step\n",
      "step:30/120 || Total Loss: 26.7583 || 1.3234s/step\n",
      "step:40/120 || Total Loss: 24.9846 || 1.4696s/step\n",
      "step:50/120 || Total Loss: 25.7564 || 1.5255s/step\n",
      "step:60/120 || Total Loss: 25.4636 || 2.4289s/step\n",
      "step:70/120 || Total Loss: 26.8181 || 1.4678s/step\n",
      "step:80/120 || Total Loss: 25.9577 || 1.9242s/step\n",
      "step:90/120 || Total Loss: 26.6623 || 1.9370s/step\n",
      "step:100/120 || Total Loss: 26.3444 || 1.5934s/step\n",
      "step:110/120 || Total Loss: 26.0586 || 1.6629s/step\n",
      "step:120/120 || Total Loss: 25.9847 || 1.7932s/step\n",
      "Finish Training.\n",
      "Total Loss: 25.7699 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:19/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 23.5791 || 1.5863s/step\n",
      "step:10/120 || Total Loss: 22.6797 || 1.4235s/step\n",
      "step:20/120 || Total Loss: 22.0037 || 1.3148s/step\n",
      "step:30/120 || Total Loss: 20.9187 || 1.4205s/step\n",
      "step:40/120 || Total Loss: 20.5161 || 1.4577s/step\n",
      "step:50/120 || Total Loss: 19.8224 || 1.8221s/step\n",
      "step:60/120 || Total Loss: 19.2765 || 1.2872s/step\n",
      "step:70/120 || Total Loss: 18.7482 || 1.3930s/step\n",
      "step:80/120 || Total Loss: 20.1944 || 1.3144s/step\n",
      "step:90/120 || Total Loss: 21.3846 || 1.9690s/step\n",
      "step:100/120 || Total Loss: 24.3318 || 2.0303s/step\n",
      "step:110/120 || Total Loss: 23.5215 || 1.4042s/step\n",
      "step:120/120 || Total Loss: 24.5402 || 1.6559s/step\n",
      "Finish Training.\n",
      "Total Loss: 24.3374 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:20/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 12.1639 || 2.4395s/step\n",
      "step:10/120 || Total Loss: 15.0753 || 1.3882s/step\n",
      "step:20/120 || Total Loss: 24.8674 || 1.3336s/step\n",
      "step:30/120 || Total Loss: 21.3581 || 2.0318s/step\n",
      "step:40/120 || Total Loss: 21.4265 || 1.4103s/step\n",
      "step:50/120 || Total Loss: 27.0476 || 1.6920s/step\n",
      "step:60/120 || Total Loss: 27.4262 || 1.5940s/step\n",
      "step:70/120 || Total Loss: 26.5973 || 1.3895s/step\n",
      "step:80/120 || Total Loss: 26.6482 || 1.4746s/step\n",
      "step:90/120 || Total Loss: 25.4573 || 1.5474s/step\n",
      "step:100/120 || Total Loss: 24.9536 || 1.5265s/step\n",
      "step:110/120 || Total Loss: 24.1316 || 1.2539s/step\n",
      "step:120/120 || Total Loss: 25.0813 || 1.8497s/step\n",
      "Finish Training.\n",
      "Total Loss: 24.8741 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:21/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 10.0247 || 1.4519s/step\n",
      "step:10/120 || Total Loss: 15.4777 || 1.4976s/step\n",
      "step:20/120 || Total Loss: 16.6885 || 1.5173s/step\n",
      "step:30/120 || Total Loss: 18.0848 || 1.6365s/step\n",
      "step:40/120 || Total Loss: 20.6149 || 1.4152s/step\n",
      "step:50/120 || Total Loss: 20.2397 || 1.3043s/step\n",
      "step:60/120 || Total Loss: 20.3828 || 1.5963s/step\n",
      "step:70/120 || Total Loss: 19.8783 || 1.8389s/step\n",
      "step:80/120 || Total Loss: 18.8014 || 1.4647s/step\n",
      "step:90/120 || Total Loss: 20.2408 || 1.5469s/step\n",
      "step:100/120 || Total Loss: 20.1551 || 1.5941s/step\n",
      "step:110/120 || Total Loss: 19.6643 || 1.2485s/step\n",
      "step:120/120 || Total Loss: 20.0062 || 1.6769s/step\n",
      "Finish Training.\n",
      "Total Loss: 19.8409 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:22/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 15.3124 || 1.4535s/step\n",
      "step:10/120 || Total Loss: 24.8094 || 1.2892s/step\n",
      "step:20/120 || Total Loss: 28.0042 || 1.9421s/step\n",
      "step:30/120 || Total Loss: 25.4075 || 1.7209s/step\n",
      "step:40/120 || Total Loss: 28.5389 || 2.5900s/step\n",
      "step:50/120 || Total Loss: 27.3846 || 2.3668s/step\n",
      "step:60/120 || Total Loss: 26.7313 || 1.7924s/step\n",
      "step:70/120 || Total Loss: 27.3972 || 1.4128s/step\n",
      "step:80/120 || Total Loss: 25.7581 || 1.4724s/step\n",
      "step:90/120 || Total Loss: 24.3593 || 1.2396s/step\n",
      "step:100/120 || Total Loss: 24.8018 || 2.7707s/step\n",
      "step:110/120 || Total Loss: 24.5346 || 2.0165s/step\n",
      "step:120/120 || Total Loss: 25.3520 || 2.2088s/step\n",
      "Finish Training.\n",
      "Total Loss: 25.1425 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:23/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 7.6582 || 1.3412s/step\n",
      "step:10/120 || Total Loss: 17.6619 || 1.3148s/step\n",
      "step:20/120 || Total Loss: 26.3189 || 1.4234s/step\n",
      "step:30/120 || Total Loss: 23.1699 || 1.6336s/step\n",
      "step:40/120 || Total Loss: 20.7074 || 1.4843s/step\n",
      "step:50/120 || Total Loss: 21.0377 || 1.4282s/step\n",
      "step:60/120 || Total Loss: 19.5190 || 1.5170s/step\n",
      "step:70/120 || Total Loss: 18.7543 || 1.5753s/step\n",
      "step:80/120 || Total Loss: 18.7863 || 1.3193s/step\n",
      "step:90/120 || Total Loss: 18.8626 || 1.3063s/step\n",
      "step:100/120 || Total Loss: 19.7254 || 1.5142s/step\n",
      "step:110/120 || Total Loss: 20.2849 || 2.1861s/step\n",
      "step:120/120 || Total Loss: 19.5576 || 1.3580s/step\n",
      "Finish Training.\n",
      "Total Loss: 19.3959 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:24/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 15.1865 || 1.3842s/step\n",
      "step:10/120 || Total Loss: 27.3325 || 1.4462s/step\n",
      "step:20/120 || Total Loss: 28.8644 || 1.2327s/step\n",
      "step:30/120 || Total Loss: 25.8245 || 1.4639s/step\n",
      "step:40/120 || Total Loss: 22.6927 || 1.5281s/step\n",
      "step:50/120 || Total Loss: 23.6483 || 1.3801s/step\n",
      "step:60/120 || Total Loss: 23.3661 || 1.2277s/step\n",
      "step:70/120 || Total Loss: 22.1484 || 1.8134s/step\n",
      "step:80/120 || Total Loss: 23.1413 || 1.6620s/step\n",
      "step:90/120 || Total Loss: 22.6475 || 2.0401s/step\n",
      "step:100/120 || Total Loss: 22.5169 || 1.2388s/step\n",
      "step:110/120 || Total Loss: 23.3328 || 1.5261s/step\n",
      "step:120/120 || Total Loss: 22.7324 || 1.6383s/step\n",
      "Finish Training.\n",
      "Total Loss: 22.5445 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:25/25\n",
      "Start Training.\n",
      "step:1/120 || Total Loss: 81.1240 || 3.1158s/step\n",
      "step:10/120 || Total Loss: 28.3821 || 1.8634s/step\n",
      "step:20/120 || Total Loss: 23.4046 || 1.5723s/step\n",
      "step:30/120 || Total Loss: 19.8581 || 1.4497s/step\n",
      "step:40/120 || Total Loss: 19.1016 || 1.5768s/step\n",
      "step:50/120 || Total Loss: 18.1552 || 1.6083s/step\n",
      "step:60/120 || Total Loss: 20.2814 || 1.5962s/step\n",
      "step:70/120 || Total Loss: 19.7628 || 1.4411s/step\n",
      "step:80/120 || Total Loss: 19.3527 || 2.2763s/step\n",
      "step:90/120 || Total Loss: 18.6439 || 1.5387s/step\n",
      "step:100/120 || Total Loss: 18.2881 || 2.0296s/step\n",
      "step:110/120 || Total Loss: 18.5721 || 1.6621s/step\n",
      "step:120/120 || Total Loss: 17.9286 || 2.0142s/step\n",
      "Finish Training.\n",
      "Total Loss: 17.7805 || Val Loss: 0.0000 \n"
     ]
    }
   ],
   "source": [
    "#------------------------------------#\n",
    "#   先冻结backbone训练\n",
    "#------------------------------------#\n",
    "lr = 1e-3\n",
    "Batch_size = 4\n",
    "Init_Epoch = 0\n",
    "Freeze_Epoch = 25\n",
    "        \n",
    "optimizer = optim.Adam(net.parameters(), lr, weight_decay=5e-4)\n",
    "if Cosine_lr:\n",
    "    lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)\n",
    "else:\n",
    "    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)\n",
    "\n",
    "gen = Generator(Batch_size, train_lines, (input_shape[0], input_shape[1])).generate(mosaic = mosaic)\n",
    "gen_val = Generator(Batch_size, val_lines, (input_shape[0], input_shape[1])).generate(mosaic = False)\n",
    "                        \n",
    "epoch_size = int(max(1, num_train//Batch_size//2.5)) if mosaic else max(1, num_train//Batch_size)\n",
    "epoch_size_val = num_val//Batch_size\n",
    "for param in model.backbone.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "best_loss = 99999999.0\n",
    "best_model_weights = copy.deepcopy(net.state_dict())\n",
    "for epoch in range(Init_Epoch, Freeze_Epoch):\n",
    "    total_loss, val_loss = fit_one_epoch(net, yolo_losses, epoch, epoch_size, epoch_size_val, gen, gen_val, \n",
    "                                         Freeze_Epoch, Cuda, optimizer, lr_scheduler)\n",
    "    if total_loss < best_loss:\n",
    "        best_loss = total_loss\n",
    "        best_model_weights = copy.deepcopy(model.state_dict())\n",
    "    with open('total_loss.csv', mode='a+') as total_loss_file:\n",
    "        total_loss_file.write(str(total_loss.item()) + '\\n')\n",
    "    #with open('val_loss.csv', mode='a+') as val_loss_file:\n",
    "    #    val_loss_file.write(str(val_loss.item()) + '\\n')\n",
    "torch.save(best_model_weights, 'model_data/yolov4_maskdetect_weights0.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:26/50\n",
      "Start Training.\n",
      "step:1/240 || Total Loss: 20.0526 || 1.0939s/step\n",
      "step:10/240 || Total Loss: 15.7005 || 0.8445s/step\n",
      "step:20/240 || Total Loss: 14.3441 || 0.8856s/step\n",
      "step:30/240 || Total Loss: 13.1000 || 0.9173s/step\n",
      "step:40/240 || Total Loss: 17.5097 || 0.9113s/step\n",
      "step:50/240 || Total Loss: 15.8164 || 0.8353s/step\n",
      "step:60/240 || Total Loss: 16.0407 || 0.8812s/step\n",
      "step:70/240 || Total Loss: 16.0285 || 0.7731s/step\n",
      "step:80/240 || Total Loss: 15.3731 || 0.7119s/step\n",
      "step:90/240 || Total Loss: 14.8034 || 0.9056s/step\n",
      "step:100/240 || Total Loss: 14.7215 || 0.7492s/step\n",
      "step:110/240 || Total Loss: 18.2945 || 1.4890s/step\n",
      "step:120/240 || Total Loss: 19.0468 || 1.0130s/step\n",
      "step:130/240 || Total Loss: 18.3268 || 0.9163s/step\n",
      "step:140/240 || Total Loss: 17.4725 || 0.8013s/step\n",
      "step:150/240 || Total Loss: 16.9751 || 0.9707s/step\n",
      "step:160/240 || Total Loss: 18.9430 || 1.2700s/step\n",
      "step:170/240 || Total Loss: 19.3538 || 0.7483s/step\n",
      "step:180/240 || Total Loss: 18.8961 || 0.8273s/step\n",
      "step:190/240 || Total Loss: 18.5310 || 0.7246s/step\n",
      "step:200/240 || Total Loss: 18.2274 || 0.9971s/step\n",
      "step:210/240 || Total Loss: 17.6345 || 0.8365s/step\n",
      "step:220/240 || Total Loss: 17.8395 || 0.9561s/step\n",
      "step:230/240 || Total Loss: 17.8692 || 0.8906s/step\n",
      "step:240/240 || Total Loss: 17.7056 || 0.9343s/step\n",
      "Finish Training.\n",
      "Total Loss: 17.6322 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:27/50\n",
      "Start Training.\n",
      "step:1/240 || Total Loss: 11.2175 || 0.9433s/step\n",
      "step:10/240 || Total Loss: 23.8770 || 0.9888s/step\n",
      "step:20/240 || Total Loss: 15.9329 || 0.8550s/step\n",
      "step:30/240 || Total Loss: 14.2239 || 1.0487s/step\n",
      "step:40/240 || Total Loss: 13.5648 || 0.9160s/step\n",
      "step:50/240 || Total Loss: 14.7075 || 1.1628s/step\n",
      "step:60/240 || Total Loss: 16.7089 || 1.2025s/step\n",
      "step:70/240 || Total Loss: 15.4734 || 0.8519s/step\n",
      "step:80/240 || Total Loss: 14.4728 || 0.9196s/step\n",
      "step:90/240 || Total Loss: 13.7875 || 0.9530s/step\n",
      "step:100/240 || Total Loss: 16.1049 || 0.9383s/step\n",
      "step:110/240 || Total Loss: 16.2294 || 1.1785s/step\n",
      "step:120/240 || Total Loss: 16.8733 || 1.0507s/step\n",
      "step:130/240 || Total Loss: 16.9658 || 0.8408s/step\n",
      "step:140/240 || Total Loss: 16.7495 || 0.9370s/step\n",
      "step:150/240 || Total Loss: 16.3994 || 0.8080s/step\n",
      "step:160/240 || Total Loss: 16.4962 || 0.8856s/step\n",
      "step:170/240 || Total Loss: 16.7014 || 0.8742s/step\n",
      "step:180/240 || Total Loss: 16.7235 || 0.8938s/step\n",
      "step:190/240 || Total Loss: 16.3319 || 1.1132s/step\n",
      "step:200/240 || Total Loss: 18.6950 || 1.2597s/step\n",
      "step:210/240 || Total Loss: 18.5147 || 0.8861s/step\n",
      "step:220/240 || Total Loss: 18.1929 || 0.8403s/step\n",
      "step:230/240 || Total Loss: 18.1343 || 1.2703s/step\n",
      "step:240/240 || Total Loss: 18.4751 || 1.3407s/step\n",
      "Finish Training.\n",
      "Total Loss: 18.3985 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:28/50\n",
      "Start Training.\n",
      "step:1/240 || Total Loss: 44.2707 || 1.2419s/step\n",
      "step:10/240 || Total Loss: 21.2810 || 1.0011s/step\n",
      "step:20/240 || Total Loss: 16.9790 || 0.8802s/step\n",
      "step:30/240 || Total Loss: 14.2008 || 0.7758s/step\n",
      "step:40/240 || Total Loss: 12.5943 || 0.8076s/step\n",
      "step:50/240 || Total Loss: 12.1946 || 0.8547s/step\n",
      "step:60/240 || Total Loss: 11.7313 || 0.8084s/step\n",
      "step:70/240 || Total Loss: 11.2481 || 0.7867s/step\n",
      "step:80/240 || Total Loss: 10.7745 || 1.1750s/step\n",
      "step:90/240 || Total Loss: 11.4450 || 0.7524s/step\n",
      "step:100/240 || Total Loss: 11.7237 || 0.8343s/step\n",
      "step:110/240 || Total Loss: 12.7256 || 1.4205s/step\n",
      "step:120/240 || Total Loss: 14.2842 || 1.4826s/step\n",
      "step:130/240 || Total Loss: 14.3441 || 1.1174s/step\n",
      "step:140/240 || Total Loss: 14.0958 || 0.9771s/step\n",
      "step:150/240 || Total Loss: 13.6787 || 0.7299s/step\n",
      "step:160/240 || Total Loss: 13.9499 || 1.0168s/step\n",
      "step:170/240 || Total Loss: 15.2666 || 0.7765s/step\n",
      "step:180/240 || Total Loss: 15.2800 || 0.8378s/step\n",
      "step:190/240 || Total Loss: 14.9819 || 0.7928s/step\n",
      "step:200/240 || Total Loss: 15.0881 || 0.9992s/step\n",
      "step:210/240 || Total Loss: 15.4635 || 0.8798s/step\n",
      "step:220/240 || Total Loss: 15.0737 || 0.7818s/step\n",
      "step:230/240 || Total Loss: 15.5608 || 0.8734s/step\n",
      "step:240/240 || Total Loss: 15.2477 || 0.9342s/step\n",
      "Finish Training.\n",
      "Total Loss: 15.1845 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:29/50\n",
      "Start Training.\n",
      "step:1/240 || Total Loss: 4.7898 || 0.8148s/step\n",
      "step:10/240 || Total Loss: 15.3025 || 1.0709s/step\n",
      "step:20/240 || Total Loss: 18.2378 || 1.0512s/step\n",
      "step:30/240 || Total Loss: 17.5283 || 1.0762s/step\n",
      "step:40/240 || Total Loss: 18.7326 || 0.8669s/step\n",
      "step:50/240 || Total Loss: 17.4883 || 0.8377s/step\n",
      "step:60/240 || Total Loss: 18.3820 || 1.2108s/step\n",
      "step:70/240 || Total Loss: 16.5491 || 0.7627s/step\n",
      "step:80/240 || Total Loss: 15.7924 || 0.8671s/step\n",
      "step:90/240 || Total Loss: 16.5230 || 0.8656s/step\n",
      "step:100/240 || Total Loss: 15.7520 || 1.3316s/step\n",
      "step:110/240 || Total Loss: 15.6677 || 0.9828s/step\n",
      "step:120/240 || Total Loss: 15.7963 || 0.7757s/step\n",
      "step:130/240 || Total Loss: 16.3399 || 1.0343s/step\n",
      "step:140/240 || Total Loss: 16.5035 || 0.8241s/step\n",
      "step:150/240 || Total Loss: 16.9634 || 0.6813s/step\n",
      "step:160/240 || Total Loss: 16.8414 || 0.9320s/step\n",
      "step:170/240 || Total Loss: 16.3628 || 0.9328s/step\n",
      "step:180/240 || Total Loss: 15.8112 || 0.8117s/step\n",
      "step:190/240 || Total Loss: 15.3651 || 0.9680s/step\n",
      "step:200/240 || Total Loss: 15.1902 || 0.8404s/step\n",
      "step:210/240 || Total Loss: 15.0174 || 0.9309s/step\n",
      "step:220/240 || Total Loss: 15.0368 || 1.0231s/step\n",
      "step:230/240 || Total Loss: 14.9939 || 1.0587s/step\n",
      "step:240/240 || Total Loss: 14.9479 || 1.0667s/step\n",
      "Finish Training.\n",
      "Total Loss: 14.8858 || Val Loss: 0.0000 \n",
      "\n",
      "----------Train one epoch.----------\n",
      "Epoch:30/50\n",
      "Start Training.\n",
      "step:1/240 || Total Loss: 32.3853 || 0.9442s/step\n",
      "step:10/240 || Total Loss: 29.0637 || 0.9271s/step\n",
      "step:20/240 || Total Loss: 19.2526 || 0.7963s/step\n",
      "step:30/240 || Total Loss: 16.3637 || 0.9230s/step\n",
      "step:40/240 || Total Loss: 16.8577 || 0.9027s/step\n",
      "step:50/240 || Total Loss: 18.2816 || 0.9163s/step\n",
      "step:60/240 || Total Loss: 17.8553 || 0.9186s/step\n",
      "step:70/240 || Total Loss: 16.4400 || 0.8657s/step\n",
      "step:80/240 || Total Loss: 15.5027 || 0.9973s/step\n",
      "step:90/240 || Total Loss: 16.7902 || 0.8760s/step\n",
      "step:100/240 || Total Loss: 16.2536 || 0.9572s/step\n",
      "step:110/240 || Total Loss: 15.7912 || 0.9779s/step\n"
     ]
    }
   ],
   "source": [
    "#------------------------------------#\n",
    "#   解冻backbone后训练\n",
    "#------------------------------------#\n",
    "lr = 1e-4\n",
    "Batch_size = 2\n",
    "Freeze_Epoch = 25\n",
    "Unfreeze_Epoch = 50\n",
    "\n",
    "optimizer = optim.Adam(net.parameters(), lr, weight_decay=5e-4)\n",
    "if Cosine_lr:\n",
    "    lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)\n",
    "else:\n",
    "    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)\n",
    "\n",
    "gen = Generator(Batch_size, train_lines, (input_shape[0], input_shape[1])).generate(mosaic = mosaic)\n",
    "gen_val = Generator(Batch_size, val_lines, (input_shape[0], input_shape[1])).generate(mosaic = False)\n",
    "                        \n",
    "epoch_size = int(max(1, num_train//Batch_size//2.5)) if mosaic else max(1, num_train//Batch_size)\n",
    "epoch_size_val = num_val//Batch_size\n",
    "for param in model.backbone.parameters():\n",
    "    param.requires_grad = True\n",
    "\n",
    "for epoch in range(Freeze_Epoch, Unfreeze_Epoch):\n",
    "    total_loss, val_loss = fit_one_epoch(net, yolo_losses, epoch, epoch_size, epoch_size_val, gen, gen_val, \n",
    "                                         Unfreeze_Epoch, Cuda, optimizer, lr_scheduler)\n",
    "    if total_loss < best_loss:\n",
    "        best_loss = total_loss\n",
    "        best_model_weights = copy.deepcopy(model.state_dict())\n",
    "    with open('total_loss.csv', mode='a+') as total_loss_file:\n",
    "        total_loss_file.write(str(total_loss.item()) + '\\n')\n",
    "    #with open('val_loss.csv', mode='a+') as val_loss_file:\n",
    "    #    val_loss_file.write(str(val_loss.item() + '\\n')\n",
    "torch.save(best_model_weights, 'model_data/yolov4_maskdetect_weights1.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
